import os
from os import path

import torch
from torch.utils.data import DataLoader, Dataset
from torchvision.datasets import ImageFolder
from torchvision import transforms
from PIL import Image
import glob



# semi=========
args = None   # 写到配置文件

img_transforms = transforms.Compose([
    transforms.ToTensor(),
    transforms.Resize(args.size),
    transforms.Normalize(
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225],
        inplace=True,
    )
])

# 增加类别
train_dataset = ImageFolder(root=args.data_root + "/train",
    transform=img_transforms)


class MVTecADTestDataset(Dataset):
    def __init__(self, root, transform, mask_transform):
        self.test_dir = path.join(root, "test")
        self.ground_truth_dir = path.join(root, "ground_truth")
        self.classes = os.listdir(self.test_dir)
        self.transform = transform
        self.mask_transform = mask_transform

        self.current_class = 0
        self.current_class_idx = 0
        self.classes_files = {
            i: os.listdir(path.join(self.test_dir, cls)) 
            for i, cls in enumerate(self.classes)
        }
    
    def __getitem__(self, index):
        if self.current_class_idx == len(self.classes_files[self.current_class]):
            self.current_class_idx = 0
            self.current_class += 1

        item_file = path.join(self.classes[self.current_class], self.classes_files[self.current_class][self.current_class_idx])
        img_file = path.join(self.test_dir, item_file)

        mask_file = item_file.replace(".png", "_mask.png")
        mask_file = path.join(self.ground_truth_dir, mask_file)
        img = Image.open(img_file)
        img = img.convert("RGB")
        img = self.transform(img)

        is_good_img = self.classes[self.current_class] == "good"
        if not is_good_img:
            mask = Image.open(mask_file)
            mask = self.mask_transform(mask)
            mask[mask != 0] = 1.
        else:
            mask = torch.zeros((1,) + img.shape[1:])

        self.current_class_idx += 1
        return img, mask, int(not is_good_img), img_file

    def __len__(self):
        return sum(len(files) for files in self.classes_files.values())


test_dataset = MVTecADTestDataset(
    root=args.data_root,
    transform=img_transforms,
    mask_transform=transforms.Compose([
        transforms.ToTensor(),
        transforms.Resize(args.size)
    ]),
)

train_dataloader = DataLoader(
    batch_size=4,
    dataset=train_dataset,
)


# Patchcore==========
#imagenet
mean_train = [0.485, 0.456, 0.406]
std_train = [0.229, 0.224, 0.225]


class MVTecDataset(Dataset):
    def __init__(self, root, transform, gt_transform, phase):
        if phase=='train':
            self.img_path = os.path.join(root, 'train')
        else:
            self.img_path = os.path.join(root, 'test')
            self.gt_path = os.path.join(root, 'ground_truth')
        self.transform = transform
        self.gt_transform = gt_transform
        # load dataset
        self.img_paths, self.gt_paths, self.labels, self.types = self.load_dataset() # self.labels => good : 0, anomaly : 1

    def load_dataset(self):

        img_tot_paths = []
        gt_tot_paths = []
        tot_labels = []
        tot_types = []

        defect_types = os.listdir(self.img_path)
        
        for defect_type in defect_types:
            if defect_type == 'good':
                img_paths = glob.glob(os.path.join(self.img_path, defect_type) + "/*.png")
                img_tot_paths.extend(img_paths)
                gt_tot_paths.extend([0]*len(img_paths))
                tot_labels.extend([0]*len(img_paths))
                tot_types.extend(['good']*len(img_paths))
            else:
                img_paths = glob.glob(os.path.join(self.img_path, defect_type) + "/*.png")
                gt_paths = glob.glob(os.path.join(self.gt_path, defect_type) + "/*.png")
                img_paths.sort()
                gt_paths.sort()
                img_tot_paths.extend(img_paths)
                gt_tot_paths.extend(gt_paths)
                tot_labels.extend([1]*len(img_paths))
                tot_types.extend([defect_type]*len(img_paths))

        assert len(img_tot_paths) == len(gt_tot_paths), "Something wrong with test and ground truth pair!"
        
        return img_tot_paths, gt_tot_paths, tot_labels, tot_types

    def __len__(self):
        return len(self.img_paths)

    def __getitem__(self, idx):
        img_path, gt, label, img_type = self.img_paths[idx], self.gt_paths[idx], self.labels[idx], self.types[idx]
        img = Image.open(img_path).convert('RGB')
        img = self.transform(img)
        if gt == 0:
            gt = torch.zeros([1, img.size()[-2], img.size()[-2]])
        else:
            gt = Image.open(gt)
            gt = self.gt_transform(gt)
        
        assert img.size()[1:] == gt.size()[1:], "image.size != gt.size !!!"
        img_cla = os.path.basename(os.path.dirname(img_path))

        return img, gt, label, img_cla + "_" + os.path.basename(img_path[:-4]), img_type


# for patchcore data
data_transforms = transforms.Compose([
                        transforms.Resize((args.load_size, args.load_size), Image.ANTIALIAS),
                        transforms.ToTensor(),
                        transforms.CenterCrop(args.input_size),
                        transforms.Normalize(mean=mean_train,
                                            std=std_train)])
gt_transforms = transforms.Compose([
                transforms.Resize((args.load_size, args.load_size)),
                transforms.ToTensor(),
                transforms.CenterCrop(args.input_size)])

inv_normalize = transforms.Normalize(mean=[-0.485/0.229, -0.456/0.224, -0.406/0.255], std=[1/0.229, 1/0.224, 1/0.255])


# KNN series========
from os.path import isdir
from pathlib import Path
import tarfile
import wget
DATASETS_PATH = Path("./data/mvtec")


def mvtec_classes():
    return [
        "bottle",
        "cable",
        "capsule",
        "carpet",
        "grid",
        "hazelnut",
        "leather",
        "metal_nut",
        "pill",
        "screw",
        # "tile",
        # "toothbrush",
        # "transistor",
        # "wood",
        # "zipper",
    ]

class MVTecDataset_knn:
    def __init__(self, cls : str, size : int = 224):
        self.cls = cls
        self.size = size
        if cls in mvtec_classes():
            self._download()
        self.train_ds = MVTecTrainDataset(cls, size)
        self.val_ds = MVTecTestDataset(cls, size)

    def _download(self):
        if not isdir(DATASETS_PATH / self.cls):
            print(f"Could not find '{self.cls}' in '{DATASETS_PATH}/'. Downloading ... ")
            url = f"ftp://guest:GU.205dldo@ftp.softronics.ch/mvtec_anomaly_detection/{self.cls}.tar.xz"
            wget.download(url)
            with tarfile.open(f"{self.cls}.tar.xz") as tar:
                tar.extractall(DATASETS_PATH)
            os.remove(f"{self.cls}.tar.xz")
            print("") # force newline

    def load(self):
        return self.train_ds, self.val_ds

class MVTecTrainDataset(ImageFolder):
    def __init__(self, cls : str, size : int):
        super().__init__(
            root=DATASETS_PATH / cls / "train",
            transform=transforms.Compose([
                transforms.Resize(256),
                transforms.CenterCrop(size),
                transforms.ToTensor(),
            ])
        )
        self.cls = cls
        self.size = size

class MVTecTestDataset(ImageFolder):
    def __init__(self, cls : str, size : int):
        super().__init__(
            root=DATASETS_PATH / cls / "test",
            transform=transforms.Compose([
                transforms.Resize(256),
                transforms.CenterCrop(size),
                transforms.ToTensor(),
            ]),
            target_transform=transforms.Compose([
                transforms.Resize(
                    256,
                    interpolation=transforms.InterpolationMode.NEAREST
                ),
                transforms.CenterCrop(size),
                transforms.ToTensor(),
            ]),
        )
        self.cls = cls
        self.size = size
            
    def __getitem__(self, index):
        path, _ = self.samples[index]
        sample = self.loader(path)
        
        if "good" in path:
            target = Image.new('L', (self.size, self.size))
            sample_class = 0
        else:
            target_path = path.replace("test", "ground_truth")
            target_path = target_path.replace(".png", "_mask.png")
            target = self.loader(target_path)
            sample_class = 1

        if self.transform is not None:
            sample = self.transform(sample)
        if self.target_transform is not None:
            target = self.target_transform(target)

        return sample, target[:1], sample_class, path


# favae, riad===========
"""
MVTec Dataset code gently borrowed from
https://github.com/byungjae89/MahalanobisAD-pytorch/blob/master/src/datasets/mvtec.py
"""
import os
import tarfile
from PIL import Image
import cv2
from tqdm import tqdm
import urllib.request
import numpy as np
import torch
from torch.utils.data import Dataset
from torchvision import transforms as T
from .transformations import centercrop_tranform


CLASS_NAMES = ['bottle', 'cable', 'capsule', 'carpet', 'grid',
               'hazelnut', 'leather', 'metal_nut', 'pill', 'screw',
               'tile', 'toothbrush', 'transistor', 'wood', 'zipper']
OBJECT = ['bottle', 'cable', 'capsule', 'hazelnut', 'metal_nut', 
            'pill', 'screw', 'toothbrush', 'transistor', 'zipper']
TEXTURE = ['carpet', 'grid', 'leather', 'tile', 'wood']

def read_img(img_path, img_color='BGR'):
    if img_color == 'gray':
        im = cv2.imread(img_path, 0)
    elif img_color == 'BGR':
        im = cv2.imread(img_path)
    elif img_color == 'RGB':
        im = cv2.imread(img_path)
        im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB)
    return im


class MVTecDataset_RE(Dataset):
    def __init__(self, data_path='../data', class_name='bottle', is_train=True,
                 resize=256, cropsize=256, transform=None, length=None, img_color='RGB'):
        # assert class_name in CLASS_NAMES, 'class_name: {}, should be in {}'.format(class_name, CLASS_NAMES)
        self.mvtec_folder_path = data_path
        self.class_name = class_name
        self.is_train = is_train
        if (isinstance(resize, list) or isinstance(resize, tuple)) and len(resize) == 2:
            self.resize = resize
        elif isinstance(resize, int):
            self.resize = (resize, resize)
        else:
            raise ValueError
        if (isinstance(cropsize, list) or isinstance(cropsize, tuple)) and len(cropsize) == 2:
            self.cropsize = cropsize
        elif isinstance(cropsize, int):
            self.cropsize = (cropsize, cropsize)
        else:
            raise ValueError


        self.length = length
        self.color = img_color

        # load dataset
        self.x, self.y, self.mask = self.load_dataset_folder()

        # set transforms
        if transform is None:
            self.transform = centercrop_tranform(img_size=self.resize, 
                                                crop_size=self.cropsize,
                                                mean=[0.485, 0.456, 0.406],
                                                std=[0.229, 0.224, 0.225])
        else:
            self.transform = transform 

    def __getitem__(self, idx):
        if idx > len(self.x) - 1:
            idx = int(idx % len(self.x))
        x, y, mask = self.x[idx], self.y[idx], self.mask[idx]
        
        x = read_img(x, self.color)
    

        if y == 0:
            mask = np.zeros([self.cropsize[0], self.cropsize[1]])
        else:
            mask = cv2.imread(mask, 0)
            
        aug = self.transform(image=x, mask=mask)
        return aug['image'], y, aug['mask']

    def __len__(self):
        if self.length is not None:
            return self.length
        else:
            return len(self.x)

    def load_dataset_folder(self):
        phase = 'train' if self.is_train else 'test'
        x, y, mask = [], [], []

        img_dir = os.path.join(self.mvtec_folder_path, self.class_name, phase)
        gt_dir = os.path.join(self.mvtec_folder_path, self.class_name, 'ground_truth')

        img_types = sorted(os.listdir(img_dir))
        for img_type in img_types:

            # load images
            img_type_dir = os.path.join(img_dir, img_type)
            if not os.path.isdir(img_type_dir):
                continue
            img_fpath_list = sorted([os.path.join(img_type_dir, f)
                                     for f in os.listdir(img_type_dir)
                                     if f.endswith('.png')])
            x.extend(img_fpath_list)

            # load gt labels
            if img_type == 'good':
                y.extend([0] * len(img_fpath_list))
                mask.extend([None] * len(img_fpath_list))
            else:
                y.extend([1] * len(img_fpath_list))
                gt_type_dir = os.path.join(gt_dir, img_type)
                img_fname_list = [os.path.splitext(os.path.basename(f))[0] for f in img_fpath_list]
                gt_fpath_list = [os.path.join(gt_type_dir, img_fname + '_mask.png')
                                 for img_fname in img_fname_list]
                mask.extend(gt_fpath_list)

        assert len(x) == len(y), 'number of x and y should be same'

        return list(x), list(y), list(mask)


# kdad======

def load_data(config):
    normal_class = config['normal_class']
    batch_size = config['batch_size']

    if config['dataset_name'] in ['mvtec']:
        data_path = '~/anomaly/data/mvtec/' + normal_class + '/train'
        mvtec_img_size = config['mvtec_img_size']

        orig_transform = transforms.Compose([
            transforms.Resize([mvtec_img_size, mvtec_img_size]),
            transforms.ToTensor()
        ])

        dataset = ImageFolder(root=data_path, transform=orig_transform)   # ImageFolder格式, 信息更少， path得继承，替换 TODO

        test_data_path = '~/anomaly/data/mvtec/' + normal_class + '/test'
        test_set = ImageFolder(root=test_data_path, transform=orig_transform)

    
    else:
        raise Exception(
            "You enter {} as dataset, which is not a valid dataset for this repository!".format(config['dataset_name']))

    train_dataloader = torch.utils.data.DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=True,
    )
    test_dataloader = torch.utils.data.DataLoader(
        test_set,
        batch_size=batch_size,
        shuffle=False,
    )

    return train_dataloader, test_dataloader


def load_localization_data(config):
    normal_class = config['normal_class']
    mvtec_img_size = config['mvtec_img_size']

    orig_transform = transforms.Compose([
        transforms.Resize([mvtec_img_size, mvtec_img_size]),
        transforms.ToTensor()
    ])

    test_data_path = '~/anomaly/data/mvtec/' + normal_class + '/test'
    test_set = ImageFolder(root=test_data_path, transform=orig_transform)
    test_dataloader = torch.utils.data.DataLoader(
        test_set,
        batch_size=512,
        shuffle=False,
    )   # here should mv the good from the test dir

    ground_data_path = '~/anomaly/data/mvtec/' + normal_class + '/ground_truth'
    ground_dataset = ImageFolder(root=ground_data_path, transform=orig_transform)
    ground_dataloader = torch.utils.data.DataLoader(
        ground_dataset,
        batch_size=512,
        num_workers=0,
        shuffle=False
    )

    x_ground = next(iter(ground_dataloader))[0].numpy()
    ground_temp = x_ground

    std_groud_temp = np.transpose(ground_temp, (0, 2, 3, 1))    # BCHW->BHWC
    x_ground = std_groud_temp

    return test_dataloader, x_ground